package org.example.config;

import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.data.document.Document;
import dev.langchain4j.data.document.loader.FileSystemDocumentLoader;
import dev.langchain4j.data.embedding.Embedding;
import dev.langchain4j.data.segment.TextSegment;
import dev.langchain4j.memory.ChatMemory;
import dev.langchain4j.memory.chat.MessageWindowChatMemory;
import dev.langchain4j.model.embedding.EmbeddingModel;
import dev.langchain4j.rag.content.Content;
import dev.langchain4j.rag.content.retriever.ContentRetriever;
import dev.langchain4j.rag.content.retriever.EmbeddingStoreContentRetriever;
import dev.langchain4j.rag.query.Query;
import dev.langchain4j.store.embedding.inmemory.InMemoryEmbeddingStore;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.annotation.Bean;
import org.springframework.context.annotation.Configuration;

import java.util.List;

@Configuration
public class QwenChatModelBean {


    @Autowired
    private EmbeddingModel embeddingModel;


    /**
     * 创建一个 QwenChatModel 对象，用于与 Qwen 模型进行交互
     * @return
     */
    @Bean
    public QwenChatModel qwenChatModel() {
        QwenChatModel qwenChatModel;
        qwenChatModel = QwenChatModel.builder()
                .apiKey("sk-bdcd99a811de4e468cd40be328a7648a")
                .modelName("qwen-max")
                .build();
        return qwenChatModel;
    }


    @Bean
    public ChatMemory chatMemory() {
        // 创建一个消息窗口的聊天记忆数据为10条
        return MessageWindowChatMemory.withMaxMessages(10);
    }



    /**
     * 创建一个内容检索器，用于从向量中检索内容
     * @return
     */
    @Bean
    public ContentRetriever contentRetriever() {
        return new ContentRetriever() {
            @Override
            public List<Content> retrieve(Query query) {
                // 1. 加载文档（整个文件作为一个 Document）
                Document document = FileSystemDocumentLoader.loadDocument(
                        "/Users/frank/JavaProjects/cloud-demo/chatAi/src/main/resources/text_document_loader.txt"
                );

                // 2. 创建内存向量存储
                InMemoryEmbeddingStore<TextSegment> embeddingStore = new InMemoryEmbeddingStore<>();

                // 3. 手动创建 TextSegment（不分割）
                TextSegment textSegment = TextSegment.from(document.text());

                // 4. 使用 EmbeddingModel 生成嵌入向量
                Embedding embedding = embeddingModel.embed(textSegment.text()).content();

                // 5. 将嵌入和文本片段添加进向量存储
                embeddingStore.add(embedding, textSegment);

                // 6. 检索与查询相关的内容
                return EmbeddingStoreContentRetriever.builder()
                        .embeddingModel(embeddingModel)
                        .embeddingStore(embeddingStore)
                        .maxResults(1)
                        .build()
                        .retrieve(query);
            }
        };
    }


}
